"""
自动清洗CSV数据实战
- 本示例演示如何使用AI自动清洗CSV数据
- 通过与AI对话，我们可以轻松完成数据清洗工作
- 本示例将演示：读取CSV文件、调用AI进行数据清洗、将结果写入新文件
"""

import os
import csv
from openai import OpenAI
from dotenv import load_dotenv

# 加载API密钥
def get_api_key():
    """从.env文件加载API密钥"""
    load_dotenv()  # 加载.env文件中的环境变量
    api_key = os.getenv("DEEPSEEK_API_KEY")
    if not api_key:
        print("警告：未找到API密钥，请确保在.env文件中设置了DEEPSEEK_API_KEY")
    return api_key

# 获取API基础URL
def get_base_url(platform="deepseek"):
    """获取不同平台的API基础URL"""
    if platform == "deepseek":
        return "https://api.deepseek.com"
    elif platform == "openai":
        return "https://api.openai.com"
    else:
        print("不支持的平台，使用默认URL")
        return "https://api.deepseek.com"

# 读取CSV文件
def read_csv_file(file_path):
    """
    读取CSV文件，返回表头和数据
    
    参数:
    - file_path: CSV文件路径
    
    返回:
    - headers: 表头列表
    - data: 数据列表（列表的列表）
    """
    print(f"正在读取文件：{file_path}")
    rows = []
    
    # 使用with语句打开文件，确保文件会被正确关闭
    with open(file_path, 'r', encoding='utf-8') as f:
        csv_reader = csv.reader(f)
        # 使用列表推导式将每一行数据添加到rows列表中
        rows = [row for row in csv_reader]
    
    # 第一行是表头
    headers = rows[0]
    # 剩余的行是数据
    data = rows[1:]
    
    print(f"读取完成，共{len(data)}条记录")
    return headers, data

# 写入CSV文件
def write_csv_file(file_path, headers, data):
    """
    将数据写入CSV文件
    
    参数:
    - file_path: 输出CSV文件路径
    - headers: 表头列表
    - data: 数据列表（列表的列表）
    """
    print(f"正在写入文件：{file_path}")
    
    # 使用with语句打开文件，确保文件会被正确关闭
    # 在Python中，`newline` 是 `open()` 函数的一个参数，用于控制换行符的处理。
    # 当使用 `csv.writer` 写入CSV文件时，若不设置 `newline=''`，在Windows系统下可能会出现多余的空行。
    # 设置 `newline=''` 可以确保 `csv.writer` 自行处理换行符，避免出现多余空行。
    with open(file_path, 'w', encoding='utf-8', newline='') as f:
        csv_writer = csv.writer(f)
        # 写入表头
        csv_writer.writerow(headers)
        # 写入数据
        csv_writer.writerows(data)
    
    print(f"文件写入完成，共{len(data)}条记录")

# 通过AI清洗数据
def ai_data_cleaning(headers, data):
    """
    通过AI清洗数据
    
    参数:
    - headers: 原始表头列表
    - data: 原始数据列表
    
    返回:
    - new_headers: 新表头列表（增加了客户等级）
    - cleaned_data: 清洗后的数据列表
    """
    print("开始调用AI进行数据清洗...")
    
    # 获取API密钥
    api_key = get_api_key()
    if not api_key:
        print("错误：未设置API密钥，无法调用AI服务")
        return headers + ["客户等级"], []  # 返回空结果
    
    # 初始化AI客户端
    client = OpenAI(api_key=api_key, base_url=get_base_url())
    
    # 将数据转换为CSV格式的字符串，便于AI理解
    csv_data = ",".join(headers) + "\n"
    for row in data:
        csv_data += ",".join(row) + "\n"
    
    # 构建提示词
    prompt = f"""
    你是一位数据清洗专家，请帮我清洗以下CSV数据，并返回清洗后的结果。

    原始数据：
    {csv_data}
    
    请按照以下要求清洗数据：
    1. 处理缺失值：
       - 如果用户ID或用户名缺失，删除该条记录
       - 如果年龄缺失，用平均年龄填充
       - 如果购买金额缺失，删除该条记录
    
    2. 处理重复值：
       - 基于用户ID去重，保留第一条记录
    
    3. 格式统一化：
       - 将所有日期格式统一为yyyy-mm-dd
    
    4. 处理异常值：
       - 过滤掉购买金额为0或异常大的值（比如>10000）
    
    5. 添加数据标注：
       - 添加一列"客户等级"，规则如下：
         * 购买金额>=1000：高价值客户
         * 购买金额>=500且<1000：重点客户
         * 购买金额<500：普通客户
    
    请只返回清洗后的CSV数据（包含表头），不要包含任何额外的解释或内容。表头应为：user_id,user_name,age,purchase_amount,purchase_date,客户等级
    """
    
    # 调用AI进行数据清洗
    response = client.chat.completions.create(
        model="deepseek-chat",  # 使用DeepSeek的聊天模型
        messages=[
            {"role": "system", "content": "你是一个专业的数据清洗助手，擅长CSV数据处理。"},
            {"role": "user", "content": prompt}
        ],
        temperature=0.1,  # 低温度，保证结果稳定性
        max_tokens=2000  # 足够的token以接收完整的清洗结果
    )
    
    # 获取AI返回的内容
    cleaned_csv = response.choices[0].message.content.strip()
    print("AI数据清洗完成")
    
    # 将AI返回的CSV文本转换回列表
    cleaned_rows = []
    for line in cleaned_csv.split('\n'):
        if line.strip():  # 忽略空行
            cleaned_rows.append(line.split(','))
    
    # 提取新的表头和数据
    if cleaned_rows:
        new_headers = cleaned_rows[0]
        cleaned_data = cleaned_rows[1:]
        print(f"清洗后的数据：{len(cleaned_data)}条记录")
        return new_headers, cleaned_data
    else:
        print("警告：AI返回的数据为空")
        return headers + ["客户等级"], []

# 通过AI分析数据
def ai_data_analysis(headers, data):
    """
    通过AI分析数据并给出建议
    
    参数:
    - headers: 表头列表
    - data: 数据列表
    
    返回:
    - analysis: AI分析结果
    """
    print("开始调用AI进行数据分析...")
    
    # 获取API密钥
    api_key = get_api_key()
    if not api_key:
        return "错误：未设置API密钥，无法调用AI服务"
    
    # 初始化AI客户端
    client = OpenAI(api_key=api_key, base_url=get_base_url())
    
    # 将数据转换为CSV格式的字符串，便于AI理解
    csv_data = ",".join(headers) + "\n"
    # map函数是Python内置函数，它会根据提供的函数对指定序列做映射。
    # 这里的map(str, row) 表示将row列表中的每个元素都转换为字符串类型。
    # 因为",".join() 方法要求传入的可迭代对象中的元素必须都是字符串类型，
    # 所以使用map函数将row中的元素都转换为字符串，以确保能正确拼接成CSV格式的字符串。
    for row in data:
        csv_data += ",".join(map(str, row)) + "\n"
    
    # 构建提示词
    prompt = f"""
    你是一位数据分析专家，请分析以下清洗后的用户购买数据，并回答问题。
    
    数据：
    {csv_data}
    
    请分析这些客户数据的特点，并给出营销建议。
    """
    
    # 调用AI进行数据分析
    response = client.chat.completions.create(
        model="deepseek-chat",  # 使用DeepSeek的聊天模型
        messages=[
            {"role": "system", "content": "你是一个专业的数据分析师，擅长从数据中发现规律和给出营销建议。"},
            {"role": "user", "content": prompt}
        ],
        temperature=0.7,  # 适中的温度，允许一些创造性
        max_tokens=1000  # 足够的token以接收详细分析
    )
    
    # 获取AI返回的内容
    analysis = response.choices[0].message.content
    print("AI数据分析完成")
    
    return analysis

# 主函数
def main():
    """主程序入口"""
    # 文件路径
    input_file = "raw_users.csv"  # 原始数据文件
    output_file = "labeled_users.csv"  # 清洗后的数据文件
    
    # 确保相对路径正确
    current_dir = os.path.dirname(os.path.abspath(__file__))
    input_path = os.path.join(current_dir, input_file)
    output_path = os.path.join(current_dir, output_file)
    
    print("=== CSV数据自动清洗示例 ===")
    print("本程序将演示如何：")
    print("1. 读取CSV文件")
    print("2. 通过AI进行数据清洗")
    print("3. 将清洗后的数据写入新文件")
    print("4. 通过AI分析清洗后的数据")
    print("==============================\n")
    
    # 步骤1：读取CSV文件
    print("\n--- 步骤1：读取CSV文件 ---")
    headers, data = read_csv_file(input_path)
    
    # 展示一些原始数据用于示例
    print("\n原始数据示例（前3条）：")
    print(f"表头: {headers}")
    for i, row in enumerate(data[:3]):
        print(f"数据{i+1}: {row}")
    
    # 步骤2：通过AI清洗数据
    print("\n--- 步骤2：通过AI清洗数据 ---")
    new_headers, cleaned_data = ai_data_cleaning(headers, data)
    
    # 展示一些清洗后的数据用于示例
    if cleaned_data:
        print("\n清洗后的数据示例（前3条）：")
        print(f"新表头: {new_headers}")
        for i, row in enumerate(cleaned_data[:3]):
            print(f"数据{i+1}: {row}")
    else:
        print("清洗后的数据为空，请检查AI服务配置")
    
    # 步骤3：将清洗后的数据写入CSV文件
    if cleaned_data:
        print("\n--- 步骤3：写入清洗后的数据 ---")
        write_csv_file(output_path, new_headers, cleaned_data)
    
    # 步骤4：通过AI分析数据
    if cleaned_data:
        print("\n--- 步骤4：通过AI分析数据 ---")
        analysis = ai_data_analysis(new_headers, cleaned_data)
        print("\nAI分析结果：")
        print(analysis)
    
    print("\n=== 程序执行完毕 ===")

# Python程序的入口点
if __name__ == "__main__":
    main()
